package com.lucifer.hawkeye.ai.nl2sql;

import com.google.common.collect.Lists;
import com.lucifer.hawkeye.ai.common.utils.SqlExtractorUtils;
import com.lucifer.hawkeye.ai.domain.nl2sql.builder.GenerateSqlBuilder;
import com.lucifer.hawkeye.ai.domain.rag.type.RagType;
import com.lucifer.hawkeye.ai.model.ChatClientFactory;
import com.lucifer.hawkeye.ai.prompt.HawkeyePrompt;
import com.lucifer.hawkeye.ai.rag.RagEngine;
import com.lucifer.hawkeye.ai.service.ExecuteSqlService;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.lang.Nullable;

import java.util.List;
import java.util.Map;
import java.util.Optional;


/**
 * @author lucifer
 * @date 2025/3/28 17:38
 */
@Slf4j
public class HawkeyeSqlEngine  implements SqlEngine{

    private final RagEngine ragEngine;

    private final ChatModel chatModel;

    private final ExecuteSqlService executeSqlService;


    @Override
    public String generateSql(String conversationId, String content,String datasource, String tableName) {
        if(StringUtils.isBlank(content)){
            throw new IllegalArgumentException("question content  cannot be null or empty");
        }
        List<Document> ddlDocumentList = ragEngine.similaritySearch(conversationId, RagType.DDL.name(),content,tableName);
        List<Document> sqlDocumentList = ragEngine.similaritySearch(conversationId,RagType.SQL.name(),content,tableName);

        GenerateSqlBuilder generateSqlBuilder = GenerateSqlBuilder.builder().content(content).ddlList(ddlDocumentList).questionSqlList(sqlDocumentList).build();
        Prompt prompt = generateSqlPrompt(generateSqlBuilder);

        String llmResponse = ChatClientFactory.buildChatClient(this.chatModel).prompt(prompt).call().content();
        assert llmResponse != null;
        if (llmResponse.contains("intermediate_sql")) {
            String intermediateSql = SqlExtractorUtils.extractSql(llmResponse);
            List<Map<String, Object>> executed = executeSql(datasource,intermediateSql);
            generateSqlBuilder.getDocumentList().add(
                    new Document(String.format("""
                        The following is a pandas DataFrame with the results of the intermediate SQL query %s:\\n%s
                        """,intermediateSql, executed.toString()
                    )));
            prompt = generateSqlPrompt(generateSqlBuilder);
            llmResponse = ChatClientFactory.buildChatClient(this.chatModel).prompt(prompt).call().content();
        }
        return SqlExtractorUtils.extractSql(llmResponse);

    }


    @Override
    public List<Map<String, Object>> executeSql(String datasource,  String sql) {
        return executeSqlService.executeSql(datasource,sql);
    }

    private Prompt generateSqlPrompt(GenerateSqlBuilder generateSqlBuilder) {
        StringBuilder ddlBuilder = new StringBuilder();
        StringBuilder sqlBuilder = new StringBuilder();
        StringBuilder questionBuilder = new StringBuilder();
        Optional.ofNullable(generateSqlBuilder.getDdlList()).orElse(Lists.newArrayList())
                .forEach(document -> ddlBuilder.append(document.getText()).append("\n"));
        SystemPromptTemplate generateSqlSystemPromptTemplate = HawkeyePrompt.generateSqlSystemPromptTemplate();
        Message systemMessage = generateSqlSystemPromptTemplate.createMessage(Map.of("ddl", ddlBuilder));
        UserMessage userMessage = new UserMessage(generateSqlBuilder.getContent());
        List<Message> messages = Lists.newArrayList(systemMessage,userMessage);
        if(!generateSqlBuilder.getQuestionSqlList().isEmpty()){
            generateSqlBuilder.getQuestionSqlList()
                    .forEach(document -> {
                        JSONObject object = JSONObject.parseObject(document.getText());
                        sqlBuilder.append(object.getString("sql")).append("\n");
                        questionBuilder.append(object.getString("question")).append("\n");
                    });
            messages.add(new AssistantMessage(sqlBuilder.toString()));
            messages.add(new UserMessage(questionBuilder.toString()));
        }
        return new Prompt(messages);
    }


    public HawkeyeSqlEngine(RagEngine ragEngine, ChatModel chatModel,ExecuteSqlService executeSqlService) {
        this.ragEngine = ragEngine;
        this.chatModel = chatModel;
        this.executeSqlService =executeSqlService;
    }

    public static Builder builder() {
        return new Builder();
    }


    public static final class Builder {

        @Nullable
        private RagEngine ragEngine;

        @Nullable
        private ChatModel chatModel;

        @Nullable
        private ExecuteSqlService executeSqlService;

        public Builder() {
        }

        public Builder ragEngine(RagEngine ragEngine) {
            this.ragEngine = ragEngine;
            return this;
        }

        public Builder chatModel(ChatModel chatModel) {
            this.chatModel = chatModel;
            return this;
        }

        public Builder executeSqlService(ExecuteSqlService executeSqlService) {
            this.executeSqlService = executeSqlService;
            return this;
        }

        public HawkeyeSqlEngine build() {
            return new HawkeyeSqlEngine(this.ragEngine,this.chatModel,this.executeSqlService);
        }
    }

}
